import cv2
import numpy as np
from sklearn.cluster import KMeans


def detect_navigation_line(image_path):
    # 读取图像
    image = cv2.imread(image_path)
    if image is None:
        print("无法读取图像文件")
        return

    # 图像预处理
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    blur = cv2.GaussianBlur(gray, (5, 5), 0)
    thresh = cv2.adaptiveThreshold(blur, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
                                   cv2.THRESH_BINARY_INV, 11, 2)

    # 形态学操作
    kernel = np.ones((3, 3), np.uint8)
    processed = cv2.morphologyEx(thresh, cv2.MORPH_OPEN, kernel, iterations=2)

    # 八连通区域分析
    num_labels, labels, stats, centroids = cv2.connectedComponentsWithStats(
        processed, connectivity=8)

    # 筛选有效区域（面积大于1000像素）
    valid_centroids = []
    for i in range(1, num_labels):
        if stats[i, cv2.CC_STAT_AREA] > 1000:
            x, y = map(int, centroids[i])
            valid_centroids.append([x, y])
            cv2.circle(image, (x, y), 5, (0, 255, 0), -1)

    if len(valid_centroids) < 2:
        print("未检测到足够的作物行")
        return image

    # K-means聚类（假设分为2类）
    kmeans = KMeans(n_clusters=2).fit(valid_centroids)
    clusters = {0: [], 1: []}
    for idx, label in enumerate(kmeans.labels_):
        clusters[label].append(valid_centroids[idx])

    # 拟合作物行直线
    def fit_line(points):
        if len(points) < 2:
            return None
        vx, vy, x0, y0 = cv2.fitLine(np.array(points, dtype=np.float32),
                                     cv2.DIST_L2, 0, 0.01, 0.01)
        cols = image.shape[1]
        return ((cols - 1, int((cols * vx - x0) * vy / vx + y0)),
                (0, int((-x0) * vy / vx + y0)))

    # 绘制作物行和导航线
    lines = []
    for cluster in clusters.values():
        if line := fit_line(cluster):
            cv2.line(image, line[0], line[1], (0, 0, 255), 2)
            lines.append(line)

    # 计算导航中线
    if len(lines) == 2:
        nav_points = []
        for i in range(2):
            nav_points.append(((lines[0][i][0] + lines[1][i][0]) // 2,
                               (lines[0][i][1] + lines[1][i][1]) // 2))
        cv2.line(image, nav_points[0], nav_points[1], (255, 0, 0), 3)

    return image






if __name__ == '__main__':
    # 加载并预处理图像
    # image = cv2.imread(r"D:\pythonProject\weed\ultralytics-main\line\img.png")
    # 使用示例
    result_image = detect_navigation_line(r"D:\pythonProject\weed\ultralytics-main\line\img.png")
    if result_image is not None:
        cv2.imshow("检测结果", result_image)
        cv2.waitKey(0)
        cv2.destroyAllWindows()
